#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import torch
import time
from model.pln import Yolov3
from utils.datasets import PLNDataset
from torch.utils.data import DataLoader
import config.yolov3_config_voc as cfg

def quick_test():
    print("快速测试优化效果")
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"设备: {device}")
    
    model = Yolov3().to(device)
    model.train()
    
    dataset = PLNDataset(anno_file_type="train", img_size=320)
    dataloader = DataLoader(dataset, batch_size=8, shuffle=False, num_workers=0)
    
    print(f"数据集大小: {len(dataset)}")
    print(f"批次数量: {len(dataloader)}")
    
    start_time = time.time()
    for i, (imgs, targets) in enumerate(dataloader):
        if i >= 5:
            break
            
        imgs = imgs.to(device)
        targets = targets.to(device)
        
        with torch.no_grad():
            predictions = model(imgs)
        
        if i == 0:
            print(f"输入形状: {imgs.shape}")
            print(f"输出形状: {predictions.shape}")
    
    end_time = time.time()
    avg_time = (end_time - start_time) / 5
    print(f"平均每个batch时间: {avg_time:.2f}秒")
    print(f"预计每个epoch时间: {avg_time * len(dataloader) / 60:.1f}分钟")
    print(f"预计总训练时间: {avg_time * len(dataloader) * 5 / 3600:.1f}小时")
    
    print("快速测试完成!")

if __name__ == "__main__":
    quick_test()
